def quick_sort(seq):
    if len(seq) <= 1:
        return seq
    else:
        pivot = seq[0]
        less = []
        greater = []
        for i in range(1, len(seq)):
            if seq[i] < pivot:
                less.append(seq[i])
            else:
                greater.append(seq[i])
    return quick_sort(less) + [pivot] + quick_sort(greater)


# 原地排序，不增加空间复杂度
# 仅遍历一次
def quick_sort_better(seq, begin, end):
    if begin < end:
        pivot = partition(seq, begin, end)
        quick_sort_better(seq, begin, pivot)
        quick_sort_better(seq, pivot + 1, end)
    return seq


def partition(seq, begin, end):  # 从小到大排序 ,左闭右开
    '''
    :param seq: 数组
    :param begin: index = 0
    :param end: index = len(seq)
    :return: index 首个元素应该在位置，index前面的所有元素均比seq[index]小，index后面的所有元素均比seq[index]大
    '''
    pivot_index = begin
    left = begin + 1  # left表示比基准值小的
    right = end - 1
    while left < right:
        if seq[left] > seq[pivot_index] and seq[right] < seq[pivot_index]:
            seq[right], seq[left] = seq[left], seq[right]
            left += 1
            right -= 1
        if seq[left] < seq[pivot_index] and left < end:
            left += 1
        if seq[right] > seq[pivot_index] and right > begin:
            right -= 1
    if seq[pivot_index] > seq[right]:  # 该位置即是排序后的正常的位置
        seq[pivot_index], seq[right] = seq[right], seq[pivot_index]
        return right
    else:
        return begin


def nth_element(seq, begin, end, rank_num):
    ranking_num = len(seq) - 1 - rank_num

    def nth_element1(seq, begin, end, ranking_num):
        pivot = partition(seq, begin, end)  # 假设pivot位置为4，rank位置为5，只需要查找右边的元素
        if ranking_num > pivot:
            return nth_element1(seq, pivot + 1, end, ranking_num)
        elif ranking_num < pivot:
            return nth_element1(seq, begin, pivot, ranking_num)
        return seq[ranking_num]

    return nth_element1(seq, begin, end, ranking_num)


def test_quick_sort():
    seq = list(range(10))
    import random
    random.shuffle(seq)
    assert quick_sort(seq) == sorted(seq)
    assert partition([2, 3, 4, 1], 0, 4) == 1
    assert partition([7, 3, 4, 5, 6], 0, 5) == 4
    assert partition([1], 0, 1) == 0
    assert partition([1, 2, 3], 0, 3) == 0
    assert partition([1, 1, 3], 0, 3) == 0
    random.shuffle(seq)
    assert quick_sort_better(seq, 0, len(seq)) == sorted(seq)


if __name__ == '__main__':
    m = [5, 1, 2, 3, 1]
    num = nth_element(m, 0, len(m), 3)
    print(num)
